package org.fujay.commons.langchain4j.server;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.output.TokenUsage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.fujay.commons.langchain4j.param.ChatParam;
import org.fujay.commons.langchain4j.param.Message;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.function.Consumer;

/**
 * @author slm
 * @description 多模态服务
 */
@Slf4j
@RequiredArgsConstructor
public class MultiService {

    private final ChatMemoryStore chatMemoryStore;
    private final ThreadPoolTaskExecutor taskExecutor;

    /**
     * 处理多模态聊天请求
     *
     * @param chatParam 聊天参数
     * @return 消息流
     */
    public Flux<Message> multimodalChat(ChatParam chatParam) {
        log.debug("开始处理多模态聊天请求: userId={}, chatId={}",
                chatParam.getUserId(), chatParam.getMemoryId());

        try {
            // 创建单播消息接收器
            Sinks.Many<Message> sink = Sinks.many().unicast().onBackpressureBuffer();
            // 创建模型
            StreamingChatModel streamingChatModel = chatParam.getStreamingChatModel();

            UserMessage userMessage = chatParam.buildMulti();

            List<ChatMessage> memory = new ArrayList<>();
            // 带记忆功能的对话
            if (ObjectUtil.isNotEmpty(chatParam.getMemoryId())) {

                ChatMemory chatMemory = createChatMemory(chatParam);
                // 添加系统消息
                if (ObjectUtil.isNotNull(chatParam.getSystemMessage())) {
                    chatMemory.add(SystemMessage.from(chatParam.getSystemMessage()));
                }
                // 添加用户消息
                chatMemory.add(userMessage);
                // 聊天记忆
                memory = chatMemory.messages();
            } else {
                memory.add(userMessage);
            }
            // 执行聊天请求
            streamingChatModel.chat(memory, createResponseHandler(chatParam, sink));

            return sink.asFlux();
        } catch (Exception e) {
            log.error("处理多模态聊天请求时发生错误: {}", e.getMessage(), e);
            return Flux.error(e);
        }
    }

    /**
     * 创建聊天记忆
     */
    private ChatMemory createChatMemory(ChatParam chatParam) {
        return MessageWindowChatMemory.builder()
                .chatMemoryStore(chatMemoryStore)
                .maxMessages(20)
                .id(chatParam.getMemoryId())
                .build();
    }

    /**
     * 创建响应处理器
     */
    private StreamingChatResponseHandler createResponseHandler(
            ChatParam chatParam, Sinks.Many<Message> sink) {

        return new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String text) {
                if (StrUtil.isNotBlank(text)) {
                    // 发送token到sink
                    sink.tryEmitNext(Message.next(
                            chatParam.getUserId(),
                            chatParam.getMemoryId(),
                            text));
                }
            }

            @Override
            public void onCompleteResponse(ChatResponse text) {

                TokenUsage tokenUsage = text.tokenUsage();
                // 发送完成消息
                Message message = Message.complete(
                        chatParam.getUserId(),
                        chatParam.getMemoryId(),
                        text.aiMessage().text(),
                        tokenUsage);
                tryEmitComplete(sink, message, chatParam.getDoOnComplete());
            }

            @Override
            public void onError(Throwable error) {
                String errorMessage;
                if (error instanceof RuntimeException && error.getCause() != null) {
                    errorMessage = error.getCause().getMessage();
                    log.error("运行时异常: {}", errorMessage, error);
                } else if (error instanceof IllegalArgumentException) {
                    errorMessage = "参数错误: " + error.getMessage();
                    log.error(errorMessage, error);
                } else {
                    errorMessage = error.getMessage();
                    log.error("未知错误: {}", errorMessage, error);
                }

                Message message = Message.error(
                        chatParam.getUserId(),
                        chatParam.getMemoryId(),
                        errorMessage);
                tryEmitComplete(sink, message, chatParam.getDoOnComplete());
            }
        };
    }

    /**
     * 完成异步回调
     *
     * @param sink     消息接收器
     * @param message  完成消息
     * @param consumer 消息消费者
     */
    private void tryEmitComplete(Sinks.Many<Message> sink, Message message, Consumer<Message> consumer) {
        sink.tryEmitNext(message);
        sink.tryEmitComplete();
        if (Objects.nonNull(consumer)) {
            taskExecutor.execute(() -> {
                try {
                    consumer.accept(message);
                } catch (Exception e) {
                    log.error("消息消费异常: {}", e.getMessage(), e);
                }
            });
        }
    }

}
